import random
import math
import networkx as nx

class GenerateNetwork:
    def __init__(self):
        pass
    
    def get_unconnect_pair(self, edges, numNodes):
        map_arr = [0] * (numNodes + 1)
        node_queue_arr = [0] * (numNodes + 1)
        start_pos = 0
        end_pos = 0
        node_num = 0
        startBfsNode = edges[0][0]
        node_queue_arr[end_pos] = startBfsNode
        end_pos += 1
        node_num = 1
        map_arr[edges[0][0]] = 1

        while start_pos < end_pos:
            tmp = node_queue_arr[start_pos]
            start_pos += 1
            for edge in edges:
                startID = edge[0]
                endID = edge[1]
                if startID == tmp and endID != tmp and map_arr[endID] == 0:
                    node_queue_arr[end_pos] = endID
                    end_pos += 1
                    node_num += 1
                    map_arr[endID] = 1
                elif endID == tmp and startID != tmp and map_arr[startID] == 0:
                    node_queue_arr[end_pos] = startID
                    end_pos += 1
                    node_num += 1
                    map_arr[startID] = 1

        if node_num == numNodes:
            return False, (0, 0)
        else:
            for nodeIndex in range(1, numNodes):
                if map_arr[nodeIndex] + map_arr[nodeIndex + 1] == 1:
                    return True, (nodeIndex, nodeIndex + 1)
            return False, (0, 0)

    def make_graph_connect(self, edges, numNodes):
        is_unconnected, unconnected_pair = self.get_unconnect_pair(edges, numNodes)
        while is_unconnected:
            edges.append(unconnected_pair)
            is_unconnected, unconnected_pair = self.get_unconnect_pair(edges, numNodes)
    
    def generate_network(self, numNodes, alpha, beta):
        generator = random.Random()
        edges = []
        for i in range(0, numNodes):
            for j in range(i, numNodes):
                if(i == j):
                    continue
                dist = math.sqrt((i - j) ** 2)
                prob = alpha * math.exp(-beta * dist)
                if generator.random() < prob:
                    edges.append((i, j))
        
        if len(edges) == 0:
            edges.append((0, 1))

        self.make_graph_connect(edges, numNodes)
        
        G = nx.Graph()
        G.add_nodes_from(range(0, numNodes))
        for tmp in edges:
            G.add_edge(tmp[0], tmp[1])
        
        return G
